#pragma region Copyright
/// <copyright>
/// Copyright (c) 2008 Ilia Shramko.
/// Released under the terms of the GNU General Public License.
/// </copyright>
#pragma endregion

#pragma once

class Security
{
public:
	typedef BOOL (WINAPI *SfcIsKeyProtectedProc)(HKEY KeyHandle, LPCWSTR SubKeyName, REGSAM KeySam); 

	static bool SfcIsKeyProtected(HKEY keyRoot, const wchar_t* keyPath)
	{
		REGSAM reg = 0;
		BOOL retValue = 1;

		HINSTANCE hlib = LoadLibrary(L"sfc.dll");
		if(hlib == NULL)
			return false;

		BOOL result = FALSE;
		SfcIsKeyProtectedProc proc = (SfcIsKeyProtectedProc)GetProcAddress(hlib, "SfcIsKeyProtected");
		if(proc != NULL)
			result = proc(keyRoot, keyPath, 0);

		FreeLibrary(hlib);
		return result != FALSE;
	}

	static HRESULT IsOwnerAdmins(HANDLE handle, SE_OBJECT_TYPE objectType)
	{
		CSid sid;
		if(!AtlGetOwnerSid(handle, SE_REGISTRY_KEY, &sid))
		{
			HRESULT hr = AtlHresultFromLastError();
			if(FAILED(hr)) return hr;
		}
		return sid == Sids::Admins() ? S_OK : S_FALSE;
	}

	static HRESULT IsAdminRunning()
	{
		bool isMember = false;
		CAccessToken procToken;
		CAccessToken impersonationToken;

		if(!procToken.GetEffectiveToken(TOKEN_READ | TOKEN_DUPLICATE) ||
			!procToken.CreateImpersonationToken(&impersonationToken) ||
			!impersonationToken.CheckTokenMembership(Sids::Admins(), &isMember))
			return E_FAIL;

		return isMember ? S_OK : S_FALSE;
	}

	static HRESULT SetOwnerSid(HKEY rootKey, const wchar_t* keyPath, CSid& sid)
	{
		CAccessToken token;
		if(!token.GetEffectiveToken(TOKEN_QUERY | TOKEN_ADJUST_PRIVILEGES))
			return E_FAIL;

		CTokenPrivileges previousPrivileges;
		if(!token.EnablePrivilege(SE_TAKE_OWNERSHIP_NAME, &previousPrivileges))
			return E_FAIL;

		CRegKey key;
		LONG error = key.Open(rootKey, keyPath, READ_CONTROL | WRITE_OWNER);
		if(error != ERROR_SUCCESS)
			return AtlHresultFromWin32(error);

		if(!AtlSetOwnerSid(key, SE_REGISTRY_KEY, sid))
		{
			HRESULT hr = AtlHresultFromLastError();
			return FAILED(hr) ? hr : E_FAIL;
		}
		return S_OK;
	}

	static HRESULT SetRequiredPermissionsToRegistryKey(HKEY rootKey, const wchar_t* keyPath)
	{
		CRegKey key;
		HRESULT hr;
		LONG error;

		// Check if the current thread has local admins SID.
		hr = IsAdminRunning();
		if(FAILED(hr)) return hr;
		if(hr == S_FALSE) return AtlHresultFromWin32(ERROR_ACCESS_DENIED);

		bool isProtected = SfcIsKeyProtected(rootKey, keyPath);
		if(isProtected)
		{
			// WRP suppresses access denied errors for protected registry keys,
			// so set ownership to local admins without further checks.
			hr = SetOwnerSid(rootKey, keyPath, Sids::Admins());
			if(FAILED(hr)) return hr;

			error = key.Open(rootKey, keyPath, WRITE_DAC);
			if(error != ERROR_SUCCESS)
				return AtlHresultFromWin32(error);
		}
		else
		{
			// WRP is not enabled for this registry keys, so we can trust to results of the key.Open call.
			error = key.Open(rootKey, keyPath, KEY_ALL_ACCESS);
			if(error != ERROR_ACCESS_DENIED)
				return AtlHresultFromWin32(error);

			// Open key to change DACL allowing FullControl access to the key.
			error = key.Open(rootKey, keyPath, WRITE_DAC);
			if(error != ERROR_SUCCESS)
			{
				//hr = IsOwnerAdmins(key, SE_REGISTRY_KEY);
				// Set ownership to local admins...
				hr = SetOwnerSid(rootKey, keyPath, Sids::Admins());
				if(FAILED(hr)) return hr;

				// ...and try again
				error = key.Open(rootKey, keyPath, WRITE_DAC);
				if(error != ERROR_SUCCESS)
					return AtlHresultFromWin32(error);
			}
		}

		// Set FullControl permissions to the registry key for the local admins.
		CDacl dacl;
		if(!AtlGetDacl(key, SE_REGISTRY_KEY, &dacl))
		{
			hr = AtlHresultFromLastError();
			return FAILED(hr) ? hr : E_FAIL;
		}

		if(!dacl.AddAllowedAce(Sids::Admins(), KEY_ALL_ACCESS, CONTAINER_INHERIT_ACE | OBJECT_INHERIT_ACE))
			return E_FAIL;

		if(!AtlSetDacl(key, SE_REGISTRY_KEY, dacl))
		{
			hr = AtlHresultFromLastError();
			return FAILED(hr) ? hr : E_FAIL;
		}

		// Try to reopen key with FullControl access to the key.
		key.Close();
		error = key.Open(rootKey, keyPath, KEY_ALL_ACCESS);
		return AtlHresultFromWin32(error);
	}
};








